package com.xkcyy.one.ai.trainable;

import com.xkcyy.one.ai.agents.OpenAIToolAgent;
import com.xkcyy.one.ai.common.ToolAgent;
import com.xkcyy.one.ai.common.tool.Tool;
import com.xkcyy.one.ai.llm.model.ChatMessage;
import com.xkcyy.one.ai.llm.model.ChatResult;
import com.xkcyy.one.dto.MessageDTO;
import com.xkcyy.one.dto.MessageRequestDTO;
import com.xkcyy.one.entity.Conversation;
import com.xkcyy.one.entity.Message;
import com.xkcyy.one.exception.BusinessException;
import com.xkcyy.one.mapper.ConversationMapper;
import com.xkcyy.one.mapper.MessageMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 工具智能体消息处理器
 * 用于处理带工具调用功能的对话
 * 
 * @author AI Assistant
 */
@Slf4j
@Service
public class ToolAgentMessageProcessor {

    private final MessageMapper messageMapper;
    private final ConversationMapper conversationMapper;
    private final ToolAgent toolAgent;
    
    @Autowired
    public ToolAgentMessageProcessor(
            MessageMapper messageMapper,
            ConversationMapper conversationMapper,
            OpenAIToolAgent toolAgent,
            List<Tool> toolExecutors) {
        this.messageMapper = messageMapper;
        this.conversationMapper = conversationMapper;
        this.toolAgent = toolAgent;
        
        // 注册所有可用的工具执行器
        if (toolExecutors != null && !toolExecutors.isEmpty()) {
            toolAgent.registerTools(toolExecutors);
            log.info("已注册 {} 个工具执行器", toolExecutors.size());
        }
    }
    
    /**
     * 保存用户消息并获取会话历史
     * 
     * @param messageRequest 消息请求
     * @return 用户消息DTO
     */
    @Transactional(rollbackFor = Exception.class)
    public MessageDTO saveUserMessageAndGetHistory(MessageRequestDTO messageRequest) {
        // 检查对话是否存在
        Conversation conversation = conversationMapper.selectById(messageRequest.getConversationId());
        if (conversation == null) {
            throw new BusinessException("对话不存在，ID: " + messageRequest.getConversationId());
        }
        
        // 创建用户消息
        Message userMessage = Message.create(Message.class);
        userMessage.setConversationId(messageRequest.getConversationId());
        userMessage.setContent(messageRequest.getContent());
        userMessage.setRole(Message.Role.USER);
        userMessage.setReasonContent(messageRequest.getReasonContent());
        
        // 保存用户消息
        messageMapper.insert(userMessage);
        
        // 更新对话时间
        conversation.setUpdateTime(LocalDateTime.now());
        conversationMapper.updateById(conversation);
        
        return MessageDTO.fromEntity(userMessage);
    }
    
    /**
     * 获取会话历史消息并转换为ChatMessage格式
     * 
     * @param conversationId 会话ID
     * @return ChatMessage列表
     */
    public List<ChatMessage> getConversationHistory(Long conversationId) {
        // 这里需要查询会话历史并转换为ChatMessage格式
        // 可以复用MessageService中的方法
        // 以下为简化实现，实际项目中应当从数据库获取
        
        List<Message> messages = messageMapper.selectList(
                new com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper<Message>()
                        .eq(Message::getConversationId, conversationId)
                        .orderByAsc(Message::getCreateTime)
        );
        
        return messages.stream()
                .map(message -> {
                    String role = Message.Role.ASSISTANT.equals(message.getRole()) ? "assistant" : "user";
                    return ChatMessage.builder()
                            .role(role)
                            .content(message.getContent())
                            .build();
                })
                .collect(Collectors.toList());
    }
    
    /**
     * 工具智能体流式聊天处理
     * 
     * @param messageRequest 消息请求
     * @return 流式聊天结果
     */
    @Transactional(rollbackFor = Exception.class)
    public Flux<ChatResult> chatSteamWithTools(MessageRequestDTO messageRequest) {
        // 保存用户消息并获取对话历史
        saveUserMessageAndGetHistory(messageRequest);
        
        // 创建助手消息实体
        Message assistantMessage = Message.create(Message.class);
        assistantMessage.setConversationId(messageRequest.getConversationId());
        assistantMessage.setContent("");
        assistantMessage.setRole(Message.Role.ASSISTANT);
        assistantMessage.setReasonContent(""); // 初始化为空，流式响应完成后可能会填充
        
        // 插入消息
        messageMapper.insert(assistantMessage);
        MessageDTO savedAssistantMessageDTO = MessageDTO.fromEntity(assistantMessage);
        
        // 用于收集完整的响应内容
        StringBuilder fullContent = new StringBuilder();
        StringBuilder reasonContent = new StringBuilder();
        
        // 获取对话历史
        List<ChatMessage> chatMessages = getConversationHistory(messageRequest.getConversationId());
        
        // 生成流式响应，并传递消息ID以便客户端可以更新
        return toolAgent.chatStream(chatMessages)
                .doOnNext(chatResult -> {
                    if (chatResult.getContent() != null) {
                        fullContent.append(chatResult.getContent());
                    }
                    // 如果有推理内容，也收集起来
                    if (chatResult.getReasoningContent() != null) {
                        reasonContent.append(chatResult.getReasoningContent());
                    }
                })
                .doOnComplete(() -> {
                    // 在流式响应完成后，更新消息内容
                    Message updateMsg = Message.create(Message.class, savedAssistantMessageDTO.getId());
                    updateMsg.setContent(fullContent.toString());
                    // 如果有收集到推理内容，也一并更新
                    if (reasonContent.length() > 0) {
                        updateMsg.setReasonContent(reasonContent.toString());
                    }
                    messageMapper.updateById(updateMsg);
                });
    }
} 